import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from itertools import count

# 全局变量
num_states = 6
states = []
state_indices = {}
transition_matrix = None
current_state = None
G = nx.DiGraph()
edge_labels = {}
pos = {}

# 生成随机合法转移矩阵（每行加起来为1）
def generate_random_transition_matrix(n):
    matrix = np.random.rand(n, n)
    row_sums = matrix.sum(axis=1, keepdims=True)
    return matrix / row_sums

# 初始化马尔可夫链
def initialize_chain(n, initial_distribution=None):
    global states, state_indices, transition_matrix, current_state, G, edge_labels, pos
    states = [chr(65 + i) for i in range(n)]  # A, B, C...
    state_indices = {state: idx for idx, state in enumerate(states)}
    transition_matrix = generate_random_transition_matrix(n)

    # 构建有向图
    G.clear()
    for state in states:
        G.add_node(state)

    edge_labels.clear()
    for i, src in enumerate(states):
        for j, dst in enumerate(states):
            prob = transition_matrix[i][j]
            if prob > 0.05:  # 只显示大于某个阈值的概率边
                G.add_edge(src, dst)
                edge_labels[(src, dst)] = f"{prob:.2f}"

    # 布局设置
    pos = nx.circular_layout(G)

    # 设置初始状态（根据初始分布选择）
    if initial_distribution is None:
        initial_distribution = np.random.dirichlet(np.ones(n), size=1)[0]  # 和为1的随机分布
    else:
        assert len(initial_distribution) == n, "Initial distribution length must match number of states"

    print("Initial State Distribution:", dict(zip(states, initial_distribution)))

    current_state = np.random.choice(states, p=initial_distribution)

# 初始化第一次状态链（带初始分布）
initial_distribution = np.random.dirichlet(np.ones(num_states), size=1)[0]
initialize_chain(num_states, initial_distribution)

# 路径记录和状态计数
path = [current_state]

# 初始化 state_count 使用初始分布
state_count = {state: (initial_distribution[i] * 100) for i, state in enumerate(states)}  # 初始权重放大便于更新
state_count[current_state] += 1  # 第一步状态计数加1

# 存储最近几次的概率分布用于判断是否收敛
history_distributions = []
stable_check_window = 10  # 最近10次分布一致认为稳定
convergence_threshold = 0.01  # 分布变化小于1%视为稳定

# 更新函数，用于动画每一帧
def update(frame):
    global current_state, num_states, path, state_count, ani

    # 自动停止逻辑：检查是否收敛
    if len(history_distributions) >= stable_check_window:
        recent_dists = history_distributions[-stable_check_window:]
        max_diff = np.max([np.abs(r - recent_dists[0]) for r in recent_dists[1:]])
        if max_diff < convergence_threshold:
            ani.event_source.stop()
            print("Distribution has converged. Animation stopped.")
            return []

    # 如果当前状态无效（比如状态数变了），重置
    if current_state not in states:
        current_state = np.random.choice(states)

    current_idx = state_indices[current_state]
    next_state = np.random.choice(states, p=transition_matrix[current_idx])
    path.append(next_state)
    current_state = next_state
    state_count[current_state] += 1  # 更新状态计数

    # 计算当前分布
    total = sum(state_count.values())
    current_distribution = np.array([state_count[state] / total for state in states])
    history_distributions.append(current_distribution)

    # 清空并重新绘制主图
    ax_main.clear()
    ax_main.set_title(f"Current State = {current_state} | States Count = {len(states)}")

    # 绘制状态转移图
    nx.draw(G, pos, with_labels=True, node_size=800, node_color='lightblue', font_size=16, ax=ax_main, arrows=True)

    # 绘制边标签（转移概率）
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, ax=ax_main, font_size=10)

    # 高亮当前状态
    nx.draw_networkx_nodes(G, pos, nodelist=[current_state], node_color='orange', node_size=900, ax=ax_main)

    # 显示路径（最近10步）
    ax_main.text(0.5, -0.1, 'Path: ' + ' → '.join(path[-10:]), transform=ax_main.transAxes,
                 ha='center', fontsize=10, color='black')

    # 更新柱状图
    ax_hist.clear()
    ax_hist.set_title("State Distribution")
    freq = [state_count[state] / total for state in states]
    bars = ax_hist.bar(states, freq, color='skyblue')
    ax_hist.set_ylim(0, 1)
    for bar in bars:
        height = bar.get_height()
        ax_hist.text(bar.get_x() + bar.get_width()/2., height,
                     f'{height:.2f}', ha='center', va='bottom')

    return []

# 控制动画暂停/继续的函数
def on_key(event):
    if event.key == 'escape':
        ani.event_source.stop()
        print("Animation stopped by pressing ESC.")

# 创建绘图对象（两个子图：左侧是图，右侧是柱状图）
fig, (ax_main, ax_hist) = plt.subplots(1, 2, figsize=(12, 6))
ani = FuncAnimation(fig, update, frames=count(), interval=100, blit=False)

# 绑定键盘事件
fig.canvas.mpl_connect('key_press_event', on_key)

plt.tight_layout()
plt.show()

# 输出最终路径
print("\nFinal Markov Chain Path:")
print(" -> ".join(path))

# 输出最后各个状态的概率分布
final_distribution = {state: count / sum(state_count.values()) for state, count in state_count.items()}
print("\nFinal State Distribution:")
for state, prob in final_distribution.items():
    print(f"{state}: {prob:.4f}")
